""" A module to handle Geoids and Datum shifts.

    :Author: 2009-2011 Roberto Vidmar

    :Copyright: 2011-2012
                Nicola Creati <ncreati@inogs.it>
                Roberto Vidmar <rvidmar@inogs.it>

    :License: MIT/X11 License (see :download:`license.txt
                               <../../license.txt>`)
"""

import interp
import numpy as np

_long   = 'i'
_double = 'd'

class DatumShift(object):
  """ A DatumShift correction class, contains the DatumShift matrix.
  """
  def __init__(self, fn):
    """ Load [delta lambda, delta phi] matrix into an array structure.

        File must be binary with the following structure:

        * 2 long integers: cols, rows
        * 4 doubles: xmin, xmax, ymin, ymax (degrees, WGS84)
        * (cols * rows) doubles: longitude data
        * (cols * rows) doubles: latitude data

      :param fn: pathname of file to load
      :type fn: string
      :raises:
    """
    self.cols = self.rows = self.xint = self.yint = 0
    self.xmin = self.xmax = self.ymin = self.ymax = None
    self._datumShiftLon = self._datumShiftLat = None

    fid = open(fn, "rb")
    # Load Header
    headerFormat = '%s,%s,%s,%s,%s,%s' % (
      _long, _long, _double, _double, _double, _double)
    header = np.fromfile(fid, dtype=headerFormat, count=1)[0]
    self.cols, self.rows, self.xmin, self.xmax, self.ymin, self.ymax = header
    data = np.fromfile(fid, dtype=_double)
    self.xint = (self.xmax - self.xmin) / self.cols
    self.yint = (self.ymax - self.ymin) / self.rows
    self._datumShiftLon = data[:data.size / 2].reshape(self.rows, self.cols)
    self._datumShiftLat = data[data.size / 2:].reshape(self.rows, self.cols)

  def __repr__(self):
    s = ''
    s += "WGS84toRoma40 is a %d x %d matrix\n" % (self.rows, self.cols)
    s += "Latitude in (%.4f, %.4f) degrees\n" % (self.ymin, self.ymax)
    s += "Longitude in (%.4f, %.4f) degrees WGS84\n" % (self.xmin, self.xmax)
    return s

  def datumShift(self):
    """ Return datum shift matrix

      :returns: datumShiftLon, datumShiftLat matrices
      :rtype: tuple
      :raises:
    """
    return self._datumShiftLon, self._datumShiftLat

  def getLLCorrection(self, xyzWGS84, method='interp2d'):
    """ Return Delta Lon/Lat correction in degrees for xyzWGS84 dataset.

        Dataset must be Lon/Lat/Z in degrees WGS84.
        Return None if at least one point is outside region

        Available methods:
          - interp2d (default)
          - thiessen (faster, no interpolation)

      :param xyzWGS84: xyz dataset (degrees WGS84)
      :type fn: numpy array
      :param method: interpolation method (optional)
      :type method: string
      :returns: loncorrection, latcorrection matices
      :rtype: tuple
      :raises:
    """
    loncorrection = latcorrection = None

    # datumShift matrix extends from xmin to xmax
    # and from ymin to ymax. Correction points are
    # computed in the center of the squares.
    #
    # ymax  +-------+-------+--...----+-------+
    #       |       |       |         |       |
    #       |   +   |  +    |         |   +   |  <---- Correction points
    #       |       |       |         |       |
    # ymax- +-------+-------+--...----+-------+
    # yint  |       |       |         |       |
    #       |       |       |         |       |
    #       .       .       .         .       .
    #       .       .       .         .       .
    #       .       .       .         .       .
    # ymin- +-------+-------+--...----+-------+
    # yint  |       |       |         |       |
    #       |   +   |   +   |         |   +   |  <---- Correction points
    #       |       |       |         |       |
    # ymin  +-------+-------+--...----+-------+
    #     xmin   xmin+              xmax-   xmax
    #            xint               xint
    #

    # Search if dataset is inside the region where matrix is defined
    if (np.min(xyzWGS84[0]) >= self.xmin and
        np.max(xyzWGS84[0]) <= self.xmax and
        np.min(xyzWGS84[1]) >= self.ymin and
        np.max(xyzWGS84[1]) <= self.ymax):
      # Ok, file fits into region, find sqare that contains point:
      # Find sqare that contains point:
      #
      # irow+1 +--------+
      #        |        |
      #        |        |
      #        |        |
      # irow   +--------+
      #       icol    icol+1
      if method == 'thiessen':
        # In this way we use Thiessen (Dirichlet/Voronoi) polygons:
        # "Value at unknown location equals value at nearest known point"
        icol = ((xyzWGS84[0] - self.xmin) / self.xint).astype(np.int32)
        irow = ((xyzWGS84[1] - self.ymin) / self.yint).astype(np.int32)
        loncorrection = - self._datumShiftLon[irow, icol]
        latcorrection = - self._datumShiftLat[irow, icol]
      else:
        xy = np.vstack((xyzWGS84[0], xyzWGS84[1])).T
        loncorrection = - interp.interp2d(self._datumShiftLon,
          ((self.xmin, self.xmax),(self.ymin, self.ymax)), xy)
        latcorrection = - interp.interp2d(self._datumShiftLat,
          ((self.xmin, self.xmax),(self.ymin, self.ymax)), xy)

    return loncorrection, latcorrection

class Geoid(object):
  """ Geoid correction class, contains the geoid matrix.
  """
  def __init__(self, fn, gtx=False):
    """ Load Geoid file into an array structure.

        - gtx == False (default)
          File must be binary with the following structure:
          2 long integers: cols, rows
          4 doubles: xmin, xmax, ymin, ymax (degrees, WGS84)
          (cols * rows) doubles: geoid data
        - gtx == True
          File must be binary with the following structure:
          2 long integers: cols, rows
          4 doubles: xmin, xmax, ymin, ymax (degrees, WGS84)
          (cols * rows) doubles: geoid data

      For GTX format: http://vdatum.noaa.gov/dev/gtx_info.html#dev_gtx_binary

      .. warning:: GTX format is **NOT** tested yet

      :param fn: pathname of file to load
      :type fn: string
      :param gtx: file has gtx format (default is False)
      :type gtx: bool
      :raises:
    """
    self.cols = self.rows = self.xint = self.yint = 0
    self.xmin = self.xmax = self.ymin = self.ymax = None
    self._geoid = None

    fid = open(fn, "rb")
    if gtx:
      # GTX format: http://vdatum.noaa.gov/dev/gtx_info.html#dev_gtx_binary
      # UNTESTED
      headerFormat = '%s,%s,%s,%s,%s,%s' % (
        _double, _double, _double, _double, _long, _long)
      # Load Header
      header = np.fromfile(fid, dtype=headerFormat, count=1)[0]
      self.ymin, self.xmin, self.yint, self.xint, self.rows, self.cols = header
      self.xmax = self.xmin + self.xint * self.cols
      self.ymax = self.ymin + self.yint * self.rows
      data = np.fromfile(fid, dtype=_double)
    else:
      headerFormat = '%s,%s,%s,%s,%s,%s' % (
        _long, _long, _double, _double, _double, _double)
      # Load Header
      header = np.fromfile(fid, dtype=headerFormat, count=1)[0]
      self.cols, self.rows, self.xmin, self.xmax, self.ymin, self.ymax = header
      data = np.fromfile(fid, dtype=_double)
      self.xint = (self.xmax - self.xmin) / self.cols
      self.yint = (self.ymax - self.ymin) / self.rows

    self._geoid = data.reshape(self.rows, self.cols)

  def __repr__(self):
    s = ''
    s += "Geoid is a %d x %d matrix\n" % (self.rows, self.cols)
    s += "Latitude in (%.4f, %.4f) degrees\n" % (self.ymin, self.ymax)
    s += "Longitude in (%.4f, %.4f) degrees WGS84\n" % (self.xmin, self.xmax)
    return s

  def geoid(self):
    """ Return the geoid matrix

      :returns: the geoid matrix
      :rtype: numpy array
      :raises:
    """
    return self._geoid

  def getCorrection(self, xyzWGS84, method='interp2d'):
    """ Return Geoid correction for xyzWGS84 dataset.

        Dataset must be Lon/Lat/Z in degrees WGS84.
        Correction must be **ADDED** to ellipsoidical height to get
        geoidical height.
        Return None if at least one point is outside region

        Available methods:
          - interp2d (default)
          - thiessen (faster, no interpolation)

      :param xyzWGS84: xyz dataset (degrees WGS84)
      :type fn: numpy array
      :param method: interpolation method (optional)
      :type method: string
      :returns: geoidcorrection matrix
      :rtype: numpy array
      :raises:
    """
    geoidcorrection = None

    # Search if dataset fits inside the geoid matrix
    if (np.min(xyzWGS84[0]) >= self.xmin and
        np.max(xyzWGS84[0]) <= self.xmax and
        np.min(xyzWGS84[1]) >= self.ymin and
        np.max(xyzWGS84[1]) <= self.ymax):
      # Ok, file fits into region, find sqare that contains point:
      #
      # irow+1 +--------+
      #        |        |
      #        |        |
      #        |        |
      # irow   +--------+
      #       icol    icol+1
      if method == 'thiessen':
        # In this way we use Thiessen (Dirichlet/Voronoi) polygons:
        # "Value at unknown location equals value at nearest known point"
        icol = ((xyzWGS84[0] - self.xmin) / self.xint).astype(np.int32)
        irow = ((xyzWGS84[1] - self.ymin) / self.yint).astype(np.int32)
        geoidcorrection = - self._geoid[irow, icol]
      else:
        xy = np.vstack((xyzWGS84[0], xyzWGS84[1])).T
        geoidcorrection = - interp.interp2d(self._geoid,
          ((self.xmin, self.xmax), (self.ymin, self.ymax)), xy)

    return geoidcorrection
